import tensorflow as tf

from codes.util.data_util import get_shape_list
from codes.util.model_util import conv2d, conv2d_transform, batch_norm, relu, lrelu, element_wise_sum, pixel_shuffer, full_connected
from codes.service.block import residual_block, discriminate_block

def generator(input_data):
    output_shape_list=get_shape_list(input_data).insert(64)
    conv_data=conv2d(input_data, output_shape_list, 3, 3, 1, 1)
    relu_data=relu(conv_data)
    old_data=relu_data

    for i in range(0,6):
        residual_data=residual_block(relu_data, output_shape_list)
        relu_data=residual_data

    conv_data = conv2d(relu_data, output_shape_list, 3, 3, 1, 1)
    norm_data = batch_norm('batch_norm', conv_data)
    relu_data = element_wise_sum(norm_data, old_data)

    output_shape_list = get_shape_list(input_data).insert(256)
    for i in range(0,2):
        conv_data = conv2d(relu_data, output_shape_list, 3, 3, 1, 1)
        shuffer_data = pixel_shuffer(conv_data)
        relu_data = relu(shuffer_data)

    output_shape_list = get_shape_list(input_data).insert(3)
    return conv2d(relu_data, output_shape_list, 3, 3, 1, 1)

def discriminator(input_data,output_shape):
    param_list = [(64,1),(64,2),(128,1),(128,2),(256,1),(256,2),(512,1),(512,2)]

    for i in range(1,9):
        if i==1:
            result_data = discriminate_block(input_data, param_list[i][0], param_list[i][1], param_list[i][1],False)
        else:
            result_data = discriminate_block(result_data, param_list[i][0], param_list[i][1], param_list[i][1], False)

    #数据拉伸
    result_data = result_data.reshape(-1,1024)
    result_data = full_connected('dense', result_data, [result_data.shape[-1],1])
    return tf.sigmoid(result_data)


